//
// Created by wujehy on 2020/5/21.
//

#include "HKDF.h"
#include "SecurityCommon.h"
#include <cmath>
#define HASH_OUTPUT_SIZE 32

#define MIN(a, b) (((a)<(b))?(a):(b))
#define MAX(a, b) (((a)>(b))?(a):(b))

using namespace GeeJoan::Common;

int HKDF::extract(std::string &prk, const std::string &salt, const std::string &input_key_material)
{
    common_hmac_sha256_base( salt, input_key_material, prk );
    return 0;
}

int HKDF::expand(std::string &output, const std::string &prk, const std::string &info, int outputlen)
{
    int iterations = (int) ceil((double) outputlen / (double) HASH_OUTPUT_SIZE );

    int result = 0;
    void *hmac_context = 0;

    std::string step_buffer;

    std::string result_buffer;

    size_t remaining_len = outputlen;
    size_t step_size = 0;

    for (int i = m_iteration_start_offset; i < iterations + m_iteration_start_offset; i++)
    {

        result = common_hmac_sha256_init( &hmac_context, prk );
        if (result < 0)break;

        if (!step_buffer.empty())
        {
            result = common_hmac_sha256_update( hmac_context, step_buffer );
            if (result < 0)break;
        }

        if (!info.empty())
        {
            result = common_hmac_sha256_update( hmac_context, info );
            if (result < 0)break;
        }

        result = common_hmac_sha256_update( hmac_context, std::to_string( i ));
        if (result < 0) break;


        result = common_hmac_sha256_final( hmac_context, step_buffer );
        if (result < 0)break;

        common_hmac_sha256_cleanup( hmac_context );

        hmac_context = 0;

        step_size = MIN( remaining_len, step_buffer.size());


        if (result_buffer.empty())
        {
            result_buffer = std::string( step_buffer.c_str(), step_size );
        } else
        {
            result_buffer = result_buffer + std::string( step_buffer.c_str(), step_size );
        }

        remaining_len -= step_size;
        result = COMMON_ERROR_CODE_SUCCESS;

    }

    if (hmac_context)
    {
        common_hmac_sha256_cleanup( hmac_context );
    }

    output = result_buffer;

    return result;
}

int HKDF::Calculate(
        std::string &outPut, const std::string &salt, const std::string &input_key_material,
        const std::string &info, int output_len)
{
    std::string prk;

    int result = COMMON_ERROR_CODE_SUCCESS;

    switch (1)
    {
        case 1:
        {
            result = this->extract( prk, salt, input_key_material );
            if (result < 0)break;
        }
        case 2:
        {
            result = this->expand( outPut, prk, info, output_len );
            if (result < 0)break;
        }
    }

    return result;
}

void HKDF::setMIterationStartOffset(int mIterationStartOffset)
{
    m_iteration_start_offset = mIterationStartOffset;
}
